package main

import (
	"context"
	"flag"
	"fmt"
	"log"

	"gitee.com/huangjianbo/mcp-server-pg/tools"
	"gitee.com/huangjianbo/mcp-server-pg/utils"
	"github.com/mark3labs/mcp-go/server"
)

var (
	Version = utils.Version
)

func newMCPServer() *server.MCPServer {
	return server.NewMCPServer(
		"mcp-server-pg",
		Version,
		server.WithToolCapabilities(true),
		server.WithLogging(),
	)
}

func addTools(s *server.MCPServer) {
	s.AddTool(tools.SQLQueryTool, tools.ExecuteSQLQueryHandler)
}

func run(transport, addr string) error {
	s := newMCPServer()
	addTools(s)

	switch transport {
	case "stdio":
		if err := server.ServeStdio(s); err != nil {
			if err == context.Canceled {
				return nil
			}
			return err
		}
	case "sse":
		srv := server.NewSSEServer(s, server.WithBaseURL("http://"+addr))
		log.Printf("SSE server listening on %s", addr)
		if err := srv.Start(addr); err != nil {
			if err == context.Canceled {
				return nil
			}
			return fmt.Errorf("failed to start sse server: %w", err)
		}
	default:
		return fmt.Errorf("invalid transport type: %s", transport)
	}

	return nil
}

func main() {
	transport := flag.String("transport", "stdio", "Transport type (stdio or sse)")
	addr := flag.String("addr", "localhost:8080", "The host and port to start the sse server on")
	flag.Parse()

	if err := run(*transport, *addr); err != nil {
		panic(err)
	}
}
